from model.fpn import FPN
from model.repvit import *
from model.fpn_head import FPNHead
from model.desc_head import Desc_head
from path import Path
from abc import ABC
import os
import time

class SegMatch(nn.Module):
    def __init__(self, configs):
        super(SegMatch, self).__init__()
        self.config = configs
        self.parameters = []

        # init backbone,neck,decoder
        self.backbone = repvit_m1_1(**self.config["RepViT_config"]).to("cuda")
        self.neck = FPN(**self.config["fpn_config"]).to("cuda")
        self.decode_head = FPNHead(**self.config["fpn_head"]).to("cuda")
        self.match_head = Desc_head().to("cuda")
        self.parameters += list(self.backbone.parameters())
        self.parameters += list(self.neck.parameters())
        self.parameters += list(self.decode_head.parameters())
        self.modules = ['neck', 'backbone',"decode_head"]

    def set_train(self):
        self.backbone.train()
        self.neck.train()
        self.decode_head.train()
        self.match_head.train()

    def set_eval(self):
        self.backbone.eval()
        self.neck.eval()
        self.decode_head.eval()
        self.match_head.eval()

    def  save_checkpoint(self,save_path):
        save_root = Path(save_path)
        model_list = ['backbone',"match_head"]
        for name in model_list:
            model_path = save_root / '{}.pth'.format(name)
            model = getattr(self, name)
            model_param = model.state_dict()
            torch.save(model_param, model_path)

    # def load_checkpoint(self, load_path):
    #     model_param = torch.load(load_path+"/backbone.pth", map_location='cuda')
    #     prefixes_to_ignore = ["encoder"]
    #     for part in ['backbone', "match_head"]:
    #         backbone_state_dict = {}
    #         head_state_dict = {}
    #         if part=="backbone":
    #             for name, param in model_param.items():
    #                     for prefix in prefixes_to_ignore:
    #                         if name.startswith(prefix):
    #                             new_name = name.split(prefix+".")[1]
    #                             backbone_state_dict[new_name] = param
    #             m = getattr(self, part)
    #             m.load_state_dict(backbone_state_dict)
    #         if part == "match_head":
    #             for name, param in model_param.items():
    #                 if not name.startswith("encoder"):
    #                     head_state_dict[name] = param
    #             m = getattr(self, part)
    #             m.load_state_dict(head_state_dict,strict=False)

    def load_checkpoint(self, load_path):
        for name in ['backbone', "match_head"]:
            model_path = load_path + '/{}.pth'.format(name)
            m = getattr(self, name)
            model_param = torch.load(model_path, map_location='cuda')
            m.load_state_dict(model_param)

    def extract(self,x):
        feature = self.backbone(x)
        xf = self.match_head(feature)
        seg_pre = self.neck(feature)
        seg_pre = self.decode_head(seg_pre)
        return {"seg_pre":seg_pre,"xf":xf}

    def forward(self, inputs):

        preds1 = self.extract(inputs['img1'])
        preds2 = self.extract(inputs['img2'])
        # preds2 = self.extract(inputs['im2'],2)

        return {"preds1":preds1,"preds2":preds2}
